import os
import pickle
import h5py
import scipy.io as sio


def confirm(prompt, yes='y', no='n'):
    t = None
    while t not in [yes, no]:
        t = input(prompt)
    if t == yes:
        return True
    elif t == no:
        return False
    else:
        raise ValueError


def confirm_file_overwrite(file_path):
    if os.path.isfile(file_path):
        if not confirm('"{}" already exists, overwrite? (y/n): '.format(file_path)):
            return False
    return True


def read_pkl_file(pkl_file_path):
    with open(pkl_file_path, 'rb') as f:
        return pickle.load(f)


def write_pkl_file(pkl_file_path, variable):
    if confirm_file_overwrite(pkl_file_path):
        with open(pkl_file_path, 'wb') as g:
            pickle.dump(variable, g, -1)


def read_mat_file(mat_file_path):
    f = sio.loadmat(mat_file_path)
    keys = [k for k in f.keys() if not k.startswith('__')]
    assert len(keys) == 1
    return f[keys[0]]


def write_mat_file(mat_file_path, variable, name, oned_as='row'):
    if confirm_file_overwrite(mat_file_path):
        sio.savemat(mat_file_path, {name: variable}, oned_as=oned_as)


class H5Reader(h5py.File):
    def __init__(self, h5_file_path, mode='r', **kwargs):
        kwargs.update({'name': h5_file_path, 'mode': mode})
        try:
            super().__init__(**kwargs)
        except:
            self._exist = False
            raise
        else:
            self._exist = True

    def __del__(self):
        if self._exist:
            self.close()


def H5Writer(h5_file_path, mode, dataset_generator, attrs_dict=None):
    if mode == 'a' or (mode == 'w' and confirm_file_overwrite(h5_file_path)):
        g = h5py.File(h5_file_path, mode=mode, libver='latest')
        try:
            if attrs_dict:
                for k in attrs_dict:
                    g.attrs[k] = attrs_dict[k]
            for name, data in dataset_generator:
                g.create_dataset(name=name, data=data)
                g.flush()
        except:
            g.close()
            raise
        else:
            g.close()
